import numpy as np
import matplotlib.pyplot as plt
import scienceplots

# 使用SCI风格
plt.style.use(["science", "no-latex"])

# 1. 生成时间序列和信号
t = np.linspace(0, 360, 10000, endpoint=False)  # 时间范围0-360
dt = t[1] - t[0]  # 采样间隔
fs = 1 / dt  # 采样频率

# 计算函数值
y = (
    20
    + 20 * np.sin((2 * np.pi * t / 360) - np.radians(110))
    + 4 * np.sin(2 * np.pi * t - 3 * np.pi / 4)
)

# 2. 傅里叶变换
n = len(y)
yf = np.fft.fft(y)
xf = np.fft.fftfreq(n, dt)

# 3. 计算振幅谱（只取正频率部分）
positive_mask = xf >= 0
xf_pos = xf[positive_mask]
yf_amp = 2.0 / n * np.abs(yf[positive_mask])
yf_amp[0] = yf_amp[0] / 2  # 校正直流分量

# 4. 创建频谱图
fig, ax = plt.subplots(figsize=(8, 5), dpi=300)

# 绘制频谱
ax.stem(
    xf_pos,
    yf_amp,
    linefmt="b-",
    markerfmt="bo",
    basefmt="r-",
    label="Amplitude Spectrum",
    # markerfacecolor="white",
    # markersize=5,
)

# 频率范围调整为0-1Hz
ax.set_xlim(0, 1)
ax.set_ylim(0, 25)

# 优化x轴刻度（0-1Hz范围内每0.2Hz一个刻度）
ax.set_xticks(np.arange(0, 1.3, 0.2))
ax.set_yticks(np.arange(0, 26, 5))

# 设置标签和标题
ax.set_xlabel("Frequency (Hz)", fontsize=12)
ax.set_ylabel("Amplitude", fontsize=12)
ax.set_title("Spectrum of Composite Sine Function (0-1Hz)", fontsize=14, pad=10)

# 添加网格和图例
ax.grid(alpha=0.5, linestyle="--")
ax.legend(frameon=True, loc="upper right")

# 优化标注位置（适应0-1Hz范围）
freq_labels = [
    (0, 20, "DC component (20)"),  # 直流分量
    (1 / 360, 20, "1/360 Hz (amplitude 20)"),  # 低频成分
    (1, 4, "1 Hz (amplitude 4)"),  # 高频成分
]

# 针对0-1Hz范围调整偏移量，避免重叠
offsets = [
    (0.05, 2),  # 直流分量：右移+上移
    (0.05, -3),  # 低频成分：右移+下移
    (-0.15, 1),  # 高频成分：左移+上移（适应右侧边界）
]

for i, (freq, amp, label) in enumerate(freq_labels):
    dx, dy = offsets[i]
    ax.annotate(
        label,
        xy=(freq, amp),
        xytext=(freq + dx, amp + dy),
        arrowprops=dict(arrowstyle="->", color="gray", lw=1),
        fontsize=10,
        bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="gray", alpha=0.7),
    )

plt.tight_layout()
plt.savefig("spectrum_0_to_1hz.png", dpi=300)
plt.show()
